'''
Author:  Andre Dozier
Date:    August 13, 2014
Purpose: Analyzes convergence characteristics of the conditional 
         gradient method used in the integrated water and power model. 
'''

import matplotlib
matplotlib.use('Agg')
import os, re
import numpy as np
import matplotlib.pyplot as plt
import access 
from modsimoutputs import getStorage

names = ['BNK_GCL','GCL_BNK','GCL_PwrRel','CHJ_PwrOut']
nodes = [2,3,4,6]
power_log_gen_cols = [7, 5, 3, 1]
power_log_target_cols = [8, 6, 4, 2]
power_log_tstep_col = 1
orig_targets_cols = [0, 1, 2, 3]
gcl_outf_i = [3,4,15,16] # [3,4,15,16,47,48,87,88]
gcl_inf_i = [1,2,14] # [1,2,14,22,70]

class CGMOutput:
    def __init__(self, cgm_log_file, power_log_file, orig_power_file, mdb_file):
        self._cgm_log(cgm_log_file)
        self._power_log(power_log_file)
        self._orig_power(orig_power_file)
        self._res_storage(mdb_file)
        self._set_len()
        self._calcDistances()
        
    def _cgm_log(self, cgm_log_file): 
        # read the entire file 
        print 'Reading ' + cgm_log_file + '...'
        f = open(cgm_log_file,'r') 
        lines = f.readlines()
        f.close()
        lines.pop(0)

        # get all timesteps
        theiter = 0
        timesteps = []
        ts = 0
        begi = 0
        endi = 0
        xs = []
        gcl_outf = [] # mass balance for GCL
        gcl_inf = [] # mass balance for GCL
        
        objs = []
        calcobjs = []
        for i in range(len(lines)): 
            currline = lines[i].split('\t')
            if int(currline[1]) < theiter or i == len(lines) - 1:
                if i == len(lines) - 1: 
                    i += 1
                endi = i
                ts += 1
                timesteps.append([begi,endi])
                begi = i
            theiter = int(currline[1])
            
            # xbar 
            xbar = np.array(currline[6].split(' '), dtype=np.float64)
            xs.append(xbar[nodes])
            gcl_outf.append(xbar[gcl_outf_i])
            gcl_inf.append(xbar[gcl_inf_i])
            if theiter == 0: 
                calcobjs.append(0)
                objs.append(0)
            else: 
                # get the variables out of the equation
                equation = currline[8]
                eqval = equation
                vars = sorted(set([int(s) for s in \
                    re.findall('q_([0-9+]*)',equation)]), reverse=True)
                eqval = eqval.replace('^','**')
                for var in vars: 
                    eqval = eqval.replace('q_{0}'.format(var), str(xbar[var-1]))

                # objective function values
                calcobjs.append(eval(eqval))
                objs.append(float(currline[2]))
            
        self.timesteps = np.array(timesteps,dtype=np.int32)
        self.decisions_iters = np.array(xs,dtype=np.float64)
        self.goal_values_iters = np.array(objs, dtype=np.float64)
        self.decisions = self._get_timestep_value(self.decisions_iters)
        self.gcl_outflow = self._get_timestep_value(gcl_outf)
        self.gcl_inflow = self._get_timestep_value(gcl_inf)
        self.goal_values_iters_calc = np.array(calcobjs, dtype=np.float64)
        indices = np.abs(self.goal_values_iters - self.goal_values_iters_calc) > np.abs(0.0001 * self.goal_values_iters_calc)
        if indices.any():
            print r"There were some goal values that were not calculated within 0.1% of IWPM's output goal values"
            print self.goal_values_iters_calc[:10]
            print self.goal_values_iters[:10]
            print np.argwhere(indices)
            plt.figure(3)
            themin = min(self.goal_values_iters.min(),self.goal_values_iters_calc.min())
            themax = max(self.goal_values_iters.max(),self.goal_values_iters_calc.max())
            plt.plot([themin,themax],[themin,themax],'k-')
            plt.plot(self.goal_values_iters, self.goal_values_iters_calc, 'x')
            plt.plot(self.goal_values_iters[indices], self.goal_values_iters_calc[indices], 'o')
            plt.ylabel('calculated')
            plt.xlabel('extracted')

    def _get_timestep_value(self, iter_values): 
        return np.array([iter_values[self.timesteps[i,1]-1] for i in range(len(self.timesteps))])
            
    def _power_log(self, power_log_file):
        print 'Reading ' + power_log_file + '...'
        power_log = np.loadtxt(power_log_file, delimiter=',', skiprows=1)
        self.generation = power_log[:,power_log_gen_cols]

    def _orig_power(self, orig_power_file):
        print 'Reading ' + orig_power_file + '...'
        orig_power = np.loadtxt(orig_power_file, delimiter=',', skiprows=1)
        self.targets = orig_power[:,orig_targets_cols]

    def _res_storage(self, mdb_file): 
        mdb = access.db(mdb_file)    
        self.gcl_storage = getStorage(mdb, 'GCL').reshape((-1))
        self.mass_balance = self.gcl_storage[:-1] + self.gcl_inflow.sum(axis=1) - self.gcl_outflow.sum(axis=1) - self.gcl_storage[1:]
        if (np.abs(self.mass_balance) > 0.1).any(): 
            print '   MASS BALANCE MISMATCH:',self.mass_balance
            print '   MEAN MASS BALANCE MISMATCH:',self.mass_balance.mean()
            
    def _set_len(self): 
        thelen = min(len(self.timesteps),len(self.decisions),len(self.generation),len(self.targets))
        self.timesteps = self.timesteps[:thelen]
        self.decisions = self.decisions[:thelen]
        self.generation = self.generation[:thelen]
        self.targets = self.targets[:thelen]

    def _calcDistances(self):
        # TODO: fix this to be the decision distances (back-calculate target decisions)
        self.distances_gen = np.linalg.norm((self.generation - self.targets),axis=1)
        self.decision_targets = self.targets / (self.generation / self.decisions)
        self.decision_targets[np.isnan(self.decision_targets)] = 0
        self.distances_decision = np.linalg.norm((self.decisions - self.decision_targets),axis=1)
        self.decision_targets_iters = np.zeros(self.decisions_iters.shape)
        for ts in range(len(self.timesteps)): 
            begi,endi = self.timesteps[ts]
            self.decision_targets_iters[begi:endi] = self.decision_targets[ts]
        self.distances_decision_iters = np.linalg.norm((self.decisions_iters - self.decision_targets_iters), axis=1)

    def getDecisionDeviation(self, timestep, byiter=False):
        if not timestep in range(len(self.timesteps)):
            return None
        
        begi,endi = self.timesteps[timestep-1]
        if byiter: 
            return self.distances_decision_iters[begi:endi]
        else: 
            return self.distances_decision[begi:endi]

    def getGenerationDeviation(self, timestep):
        if not timestep in range(len(self.timesteps)):
            return None
        
        begi,endi = self.timesteps[timestep-1]
        return self.distances_decision[begi:endi]

    def getGoalValues(self, timestep):
        if not timestep in range(len(self.timesteps)):
            return None

        begi,endi = self.timesteps[timestep-1]
        return self.goal_values_iters[begi:endi]

if __name__ == '__main__':
    # start by closing all figures
    plt.close('all')
    plt.ion()

    # targets 
    targets = [0., 408., 9760.3, 11479.3]

    # the file with the iteration info in it
    basename = os.path.join('..','GCL_TwoBus','Debug','debug_H_1e+6_M_1500_S_0')
    cgmfile = basename + '_log_cgm.txt'
    powerlogfile = basename + '_power.csv'
    powerhistfile = 'original_targets.csv'
    mdbfile = basename + 'OUTPUT.mdb'
    cgmlog = CGMOutput(cgmfile,powerlogfile,powerhistfile,mdbfile)

    # let user know number of timesteps
    n = len(cgmlog.timesteps)
    print 'Number of timesteps:',n

    # open up figures early 
    plt.figure(1)
    plt.figure(2)

    # parameters
    while True:
                            
        # which timestep to analyze?
        timestep = raw_input('Timestep to analyze (q to quit): ')
        if timestep[:1].lower() == 'q':
            break
        else:
            try:
                timestep = int(timestep)
            except:
                print '\n   YOU MUST PROVIDE AN INTEGER!\n'
                continue

        # collect only data at the first timestep 
        dists = cgmlog.getDecisionDeviation(timestep,byiter=True)
        objs = cgmlog.getGoalValues(timestep)
        if dists == None or objs == None:
            print '\n    YOU MUST PROVIDE A NUMBER BETWEEN 1 AND {0}.'.format(n)
            continue
        
        # distances       
        plt.figure(1)
        plt.plot(dists[1:])
        plt.ylabel('Euclidean distances')
        plt.xlabel('Iteration')
              
        # objective function values
        plt.figure(2)
        plt.plot(objs[1:])
        plt.ylabel('Objective function value')
        plt.xlabel('Iteration')

